module aru_reduce_max_min_4_to_1 (
    input  logic  clk,
    input  logic  rst_n,
    input  logic  en,
    input  logic  is_max,
    input  logic  psum_en,
    input  logic  end_reduce,
    input  logic  dat_in_vld [3:0],
    input  bf16_t dat_in     [3:0],
    output bf16_t dat_out,
    output logic  dat_out_vld
);

    localparam int DWIDTH = 16;
    localparam int EWIDTH = 8;
    localparam int MWIDTH = 7;
    localparam int BIAS = 127;
    localparam int INPUTS = 4;

    localparam logic [EWIDTH-1:0] EXP_ZEROS = '0;
    localparam logic [EWIDTH-1:0] EXP_ONES = '1;
    localparam logic [MWIDTH-1:0] MANT_ZEROS = '0;
    localparam logic [15:0] POS_INF = 16'h7F80;
    localparam logic [15:0] NEG_INF = 16'hFF80;
    localparam logic [15:0] NAN = 16'h7FC0;

    typedef struct packed {
        logic is_inf;
        logic is_nan;
        logic is_zero;
    } special_flags_t;

    //==========================================================================
    // Stage 1: 输入寄存 + 特殊值检测
    //==========================================================================
    bf16_t          operands_s1     [INPUTS];
    special_flags_t special_flags_s1[INPUTS];
    logic           is_max_s1;
    logic           psum_en_s1;
    logic           end_reduce_s1;

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            is_max_s1     <= '0;
            psum_en_s1    <= '0;
            end_reduce_s1 <= '0;
            for (int i = 0; i < INPUTS; i++) begin
                operands_s1[i]      <= '0;
                special_flags_s1[i] <= '0;
            end
        end else if (en) begin
            is_max_s1     <= is_max;
            psum_en_s1    <= psum_en;
            end_reduce_s1 <= end_reduce;

            for (int i = 0; i < INPUTS; i++) begin
                if (dat_in_vld[i]) begin
                    operands_s1[i]              <= dat_in[i];
                    special_flags_s1[i].is_inf  <= (&dat_in[i].exp) & (~|dat_in[i].mant);
                    special_flags_s1[i].is_nan  <= (&dat_in[i].exp) & (|dat_in[i].mant);
                    special_flags_s1[i].is_zero <= (~|dat_in[i].exp) & (~|dat_in[i].mant);
                end else begin
                    // 无效输入设为极值
                    operands_s1[i]              <= is_max ? NEG_INF : POS_INF;
                    special_flags_s1[i].is_inf  <= 1'b1;
                    special_flags_s1[i].is_nan  <= 1'b0;
                    special_flags_s1[i].is_zero <= 1'b0;
                end
            end
        end
    end

    //==========================================================================
    // 比较函数
    //==========================================================================
    function automatic bf16_t compare_bf16(input bf16_t a, input bf16_t b, input logic is_max_op,
                                           input special_flags_t a_flags, input special_flags_t b_flags);
        bf16_t result;

        if (a_flags.is_nan || b_flags.is_nan) begin
            result = NAN;
        end else if (a_flags.is_zero && b_flags.is_zero) begin
            result = is_max_op ? {1'b0, EXP_ZEROS, MANT_ZEROS} : {1'b1, EXP_ZEROS, MANT_ZEROS};
        end else begin
            logic a_larger;

            if (a.sign != b.sign) begin
                a_larger = ~a.sign;
            end else begin
                logic exp_larger, mant_larger, exp_equal;
                exp_equal   = (a.exp == b.exp);
                exp_larger  = (a.exp > b.exp);
                mant_larger = (a.mant > b.mant);

                if (a.sign == 1'b0) begin
                    a_larger = exp_larger | (exp_equal & mant_larger);
                end else begin
                    a_larger = ~(exp_larger | (exp_equal & mant_larger));
                end
            end

            result = (is_max_op == a_larger) ? a : b;
        end

        return result;
    endfunction

    function automatic special_flags_t get_special_flags(input bf16_t data);
        special_flags_t flags;
        flags.is_nan  = (&data.exp) & (|data.mant);
        flags.is_inf  = (&data.exp) & (~|data.mant);
        flags.is_zero = (~|data.exp) & (~|data.mant);
        return flags;
    endfunction

    //==========================================================================
    // Stage 2: 4输入比较（树形结构）+ 与历史值比较
    //==========================================================================

    // 历史值寄存器
    bf16_t history_result;
    logic  history_valid;

    // Stage 2 组合逻辑信号
    bf16_t compare_pair_0, compare_pair_1;
    bf16_t current_result;
    bf16_t final_result;

    // Stage 2 寄存器
    bf16_t final_result_s2;
    logic  end_reduce_s2;

    // 第一层：两两比较
    always_comb begin
        compare_pair_0 =
            compare_bf16(operands_s1[0], operands_s1[1], is_max_s1, special_flags_s1[0], special_flags_s1[1]);
        compare_pair_1 =
            compare_bf16(operands_s1[2], operands_s1[3], is_max_s1, special_flags_s1[2], special_flags_s1[3]);
    end

    // 第二层：合并结果
    always_comb begin
        current_result = compare_bf16(
            compare_pair_0,
            compare_pair_1,
            is_max_s1,
            get_special_flags(
                compare_pair_0
            ),
            get_special_flags(
                compare_pair_1)
        );
    end

    // 与历史值比较（psum_en启用时）
    always_comb begin
        if (psum_en_s1 && history_valid) begin
            final_result = compare_bf16(
                current_result,
                history_result,
                is_max_s1,
                get_special_flags(
                    current_result
                ),
                get_special_flags(
                    history_result)
            );
        end else begin
            final_result = current_result;
        end
    end

    // 更新历史值寄存器
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            history_result <= '0;
            history_valid  <= '0;
        end else if (end_reduce_s1) begin
            // end_reduce时清除历史值
            history_result <= '0;
            history_valid  <= '0;
        end else if (psum_en_s1 && en) begin
            // psum_en时更新历史值
            history_result <= final_result;
            history_valid  <= 1'b1;
        end
    end

    // Stage 2 寄存器
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            final_result_s2 <= '0;
            end_reduce_s2   <= '0;
        end else if (en) begin
            final_result_s2 <= final_result;
            end_reduce_s2   <= end_reduce_s1;
        end
    end

    //==========================================================================
    // Stage 3: 流水线寄存器
    //==========================================================================
    bf16_t final_result_s3;
    logic  end_reduce_s3;

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            final_result_s3 <= '0;
            end_reduce_s3   <= '0;
        end else if (en) begin
            final_result_s3 <= final_result_s2;
            end_reduce_s3   <= end_reduce_s2;
        end
    end

    //==========================================================================
    // Stage 4: 流水线寄存器
    //==========================================================================
    bf16_t final_result_s4;
    logic  end_reduce_s4;

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            final_result_s4 <= '0;
            end_reduce_s4   <= '0;
        end else if (en) begin
            final_result_s4 <= final_result_s3;
            end_reduce_s4   <= end_reduce_s3;
        end
    end

    //==========================================================================
    // Stage 5: 输出
    //==========================================================================
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            dat_out     <= '0;
            dat_out_vld <= '0;
        end else if (end_reduce_s4 && en) begin
            dat_out     <= final_result_s4;
            dat_out_vld <= 1'b1;
        end else begin
            dat_out_vld <= '0;
        end
    end

endmodule
